#!/usr/bin/env python3
import argparse
import os
import sys
import json
from time import sleep

import grpc

sys.path.append(
    os.path.join(os.path.dirname(os.path.abspath(__file__)),
                 '../../utils/'))
import p4runtime_lib.bmv2
import p4runtime_lib.helper
from p4runtime_lib.error_utils import printGrpcError
from p4runtime_lib.switch import ShutdownAllSwitchConnections

def configureSwitch(p4info_helper, switch, switch_config):
    """Configure a single switch with forwarding rules"""
    print(f"Configuring {switch.name}")

    # 安装基本的转发规则
    for rule in switch_config['forwarding_rules']:
        table_entry = p4info_helper.buildTableEntry(
            table_name="MyIngress.ipv4_lpm",
            match_fields={
                "hdr.ipv4.dstAddr": (rule['dst_ip'], 32)
            },
            action_name="MyIngress.ipv4_forward",
            action_params={
                "dstAddr": rule['dst_mac'],
                "port": rule['port']
            })
        switch.WriteTableEntry(table_entry)
        print(f"  Installed route: {rule['dst_ip']} -> port {rule['port']}")

def main(p4info_file_path, bmv2_file_path):
    p4info_helper = p4runtime_lib.helper.P4InfoHelper(p4info_file_path)

    # 配置数据 - 基于拓扑
    switch_configs = {
        's1': {
            'forwarding_rules': [
                {'dst_ip': '10.0.2.2', 'dst_mac': '08:00:00:00:02:22', 'port': 2},
                {'dst_ip': '10.0.2.22', 'dst_mac': '08:00:00:00:02:22', 'port': 2},
                {'dst_ip': '10.0.3.3', 'dst_mac': '08:00:00:00:03:33', 'port': 3}
            ]
        },
        's2': {
            'forwarding_rules': [
                {'dst_ip': '10.0.1.1', 'dst_mac': '08:00:00:00:01:11', 'port': 2},
                {'dst_ip': '10.0.1.11', 'dst_mac': '08:00:00:00:01:11', 'port': 2},
                {'dst_ip': '10.0.3.3', 'dst_mac': '08:00:00:00:03:33', 'port': 3}
            ]
        },
        's3': {
            'forwarding_rules': [
                {'dst_ip': '10.0.1.1', 'dst_mac': '08:00:00:00:01:11', 'port': 2},
                {'dst_ip': '10.0.1.11', 'dst_mac': '08:00:00:00:01:11', 'port': 2},
                {'dst_ip': '10.0.2.2', 'dst_mac': '08:00:00:00:02:22', 'port': 3},
                {'dst_ip': '10.0.2.22', 'dst_mac': '08:00:00:00:02:22', 'port': 3}
            ]
        }
    }

    try:
        # 确保logs目录存在且有写权限
        os.makedirs('logs', exist_ok=True)

        # 连接到交换机
        s1 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s1', address='127.0.0.1:50051', device_id=0,
            proto_dump_file='logs/s1-p4runtime-requests.txt')
        s2 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s2', address='127.0.0.1:50052', device_id=1,
            proto_dump_file='logs/s2-p4runtime-requests.txt')
        s3 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s3', address='127.0.0.1:50053', device_id=2,
            proto_dump_file='logs/s3-p4runtime-requests.txt')

        # 主仲裁更新
        s1.MasterArbitrationUpdate()
        s2.MasterArbitrationUpdate()
        s3.MasterArbitrationUpdate()

        # 安装P4程序
        s1.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        s2.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        s3.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        print("P4 programs installed on all switches")

        # 配置每个交换机
        configureSwitch(p4info_helper, s1, switch_configs['s1'])
        configureSwitch(p4info_helper, s2, switch_configs['s2'])
        configureSwitch(p4info_helper, s3, switch_configs['s3'])

        print("\nECN configuration complete!")
        print("The switches are now ready for ECN testing.")
        print("ECN threshold is set in the P4 code (ECN_THRESHOLD constant).")
        print("To change the threshold, modify ecn.p4 and recompile.")

        # 保持运行
        while True:
            sleep(10)

    except KeyboardInterrupt:
        print("\nShutting down...")
    except Exception as e:
        print(f"Error: {e}")

    ShutdownAllSwitchConnections()

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='ECN Controller')
    parser.add_argument('--p4info', help='p4info file',
                        type=str, default='./build/ecn.p4.p4info.txt')
    parser.add_argument('--bmv2-json', help='BMv2 JSON file',
                        type=str, default='./build/ecn.json')
    args = parser.parse_args()

    if not os.path.exists(args.p4info):
        print(f"P4Info file not found: {args.p4info}")
        sys.exit(1)
    if not os.path.exists(args.bmv2_json):
        print(f"BMv2 JSON file not found: {args.bmv2_json}")
        sys.exit(1)

    main(args.p4info, args.bmv2_json)
